# -*- coding: utf-8 -*-
import numpy as np
import warnings

def partition(X, rates=[0.7, 0.3]):
    assert(1 == sum(rates))
    begin = 0
    end = 0
    _X = []
    for r in rates:
        begin = end
        end += round(len(X) * r)
        _X.append(X[begin:end])
    return _X

def partition_pair(X, Y, rates=[0.7, 0.3]):
    _X = partition(X, rates)
    _Y = partition(Y, rates)
    return _X,_Y


def partition_n_folds(X, n_folds=3):
    X_split = []

    fold_size = int(len(X) / n_folds)

    begin = 0
    for i in range(n_folds):
        X_split.append(X[begin:begin + fold_size])
        begin += fold_size

    return X_split


def partition_pair_n_folds(X, Y, n_folds=3):
    _X = partition_n_folds(X, n_folds)
    _Y = partition_n_folds(Y, n_folds)
    return _X, _Y

def partition_data(method,Xs,Ys,cut_class=True,least_class_length=10):
    warnings.warn("This function was deprecated"
                  "Using PartitionData instead", DeprecationWarning)
    # =============================================
    # Aim: Cut data to training set and test set 
    # Main Inputs: Xs,Ys
    #     Two method: 'train_test' and 'train_dev_test'
    # Outputs: partioned data
    # =============================================
    
    dict_Xs = {}
    dict_Ys = {}
    Xs = np.array(Xs)
    Ys = np.array(Ys)
    np.random.seed(42)
    
    if cut_class == False: # partition data regardless of category
        permutation = np.random.permutation(len(Ys))
        Xs = Xs[permutation] # random order
        Ys = Ys[permutation] # random order
        if method == 'train_test':
            stop = int(len(Ys) * 0.7)
            dict_Xs['train'] = Xs[0:stop]
            dict_Ys['train'] = Ys[0:stop]
            dict_Xs['test'] = Xs[stop:]
            dict_Ys['test'] = Ys[stop:]
        elif method == 'train_dev_test':
            stop1 = int(len(Ys) * 0.6)
            stop2 = int(len(Ys) * 0.8)
            dict_Xs['train'] = Xs[0:stop1]
            dict_Ys['train'] = Ys[0:stop1]
            dict_Xs['dev'] = Xs[stop1:stop2]
            dict_Ys['dev'] = Ys[stop1:stop2]
            dict_Xs['test'] = Xs[stop2:]
            dict_Ys['test'] = Ys[stop2:]
    else: # partition data according to each category
        dict_Xs['train'] = []
        dict_Ys['train'] = []
        dict_Xs['test'] = []
        dict_Ys['test'] = []
        if method == 'train_dev_test':
            dict_Xs['dev'] = []
            dict_Ys['dev'] = []
            
        dict_diff_ctgr = {}
        for i in range(len(Xs)):
            X = Xs[i]
            Y = Ys[i]
            if Y not in dict_diff_ctgr:
                dict_diff_ctgr[Y] = []
            dict_diff_ctgr[Y].append(X)
            
        seeds = np.random.permutation(len(dict_diff_ctgr))
        cnt = 0
        for Y in dict_diff_ctgr: # for each category
            np.random.seed(seeds[cnt]+5)
            ctgr_Xs = np.array(dict_diff_ctgr[Y])
            len_Y = len(ctgr_Xs)
            if len_Y < least_class_length:
                print ("# %s : %d, less than %d. Delete." %(Y,len_Y,least_class_length))
            else:
                permutation = np.random.permutation(len_Y)
                ctgr_Xs = ctgr_Xs[permutation] # random order in each category
                ctgr_Ys = np.full(len_Y,Y,dtype='|S%d'%(len(Y))) # random order in each category
                if method == 'train_test':
                    stop = int(len_Y * 0.7)
                    dict_Xs['train'].extend(ctgr_Xs[0:stop])
                    dict_Ys['train'].extend(ctgr_Ys[0:stop])
                    dict_Xs['test'].extend(ctgr_Xs[stop:])
                    dict_Ys['test'].extend(ctgr_Ys[stop:])
                elif method == 'train_dev_test':
                    stop1 = int(len_Y * 0.6)
                    stop2 = int(len_Y * 0.8)
                    dict_Xs['train'].extend(ctgr_Xs[0:stop1])
                    dict_Ys['train'].extend(ctgr_Ys[0:stop1])
                    dict_Xs['dev'].extend(ctgr_Xs[stop1:stop2])
                    dict_Ys['dev'].extend(ctgr_Ys[stop1:stop2])
                    dict_Xs['test'].extend(ctgr_Xs[stop2:])
                    dict_Ys['test'].extend(ctgr_Ys[stop2:])
                    
    return dict_Xs,dict_Ys

# ========================================================================================================    
# ========================================================================================================

def CheckDuplicateData(X,Y,X_add,Y_add):
    X_add_remain = []
    Y_add_remain = []
    print ('Check duplication data: data in test set should not be in training set too...')
    for i in range(len(X_add)):
        if X_add[i].decode('utf-8') not in X:
            X_add_remain.append(X_add[i])
            Y_add_remain.append(Y_add[i])
        if i % 10000 == 0:
            print ('processing: %d / %d ' % (i,len(X_add)))
    return X_add_remain,Y_add_remain